#!/usr/bin/env python3
import argparse
import json
import os

import requests

DOCKER_USER_AGENT = r"docker/20.10.24+dfsg1 go/go1.19.8 git-commit/5d6db84 kernel/6.1.0-32-amd64 os/linux arch/amd64 UpstreamClient(Docker-Client/20.10.24+dfsg1 \(linux\))"


def parse_repo_tag(repo_tag):
    if ":" not in repo_tag:
        repo_tag += ":latest"
    repository, tag = repo_tag.split(":")
    first_part = repository.split("/")[0]
    if "." not in first_part:
        host = "registry-1.docker.io"
        img = repository
    else:
        host = first_part
        img = "/".join(repository.split("/")[1:])

    if "/" not in img:
        img = "library/" + img
    return host, img, tag


def get_auth_url(auth_header, img):
    # Parse Www-Authenticate header
    auth_header = auth_header.replace("Bearer ", "")
    auth_headers = auth_header.split(",")
    realm = None
    scope = None
    service = None
    for h in auth_headers:
        k, v = h.split("=")
        if k == "realm":
            realm = v.strip('"')
        elif k == "service":
            service = v.strip('"')
        elif k == "scope":
            scope = v

    scope = f"repository:{img}:pull"

    print("get auth token...")
    auth_url = (
        f"{realm}?service={service}&scope={scope}"
        if service
        else f"{realm}?scope={scope}"
    )
    return auth_url


def docker_pull(repo_tag):
    save_path = repo_tag.replace("/", "_").replace(":", "_")
    blobs_path = f"{save_path}/blobs/sha256"
    os.makedirs(blobs_path, exist_ok=True)

    host, img, tag = parse_repo_tag(repo_tag)

    print(f"Host: {host}")
    print(f"Image: {img}")
    print(f"Tag: {tag}")

    token = None
    session = requests.Session()
    session.headers.update({"User-Agent": DOCKER_USER_AGENT})

    print("|> check /v2/ endpoint and auth")
    resp = session.head(f"https://{host}/v2/")
    http_code = resp.status_code

    if resp.headers.get("Docker-Distribution-Api-Version") != "registry/2.0":
        print("Docker-Distribution-Api-Version should be registry/2.0")
        return

    if http_code == 200:
        print("/v2/ request success")
    elif http_code == 401:
        print("/v2/ requires auth")
        auth_header = resp.headers.get("Www-Authenticate")
        if not auth_header:
            print("error: no Www-Authenticate header")
            return

        auth_url = get_auth_url(auth_header, img)
        resp = session.get(auth_url)
        if resp.status_code != 200:
            print(f"error: failed to get auth token {resp.status_code}")
            return
        token_data = resp.json()
        token = token_data.get("token")
        if not token:
            print("error: token not found in auth response")
            return
        session.headers.update({"Authorization": f"Bearer {token}"})
    elif http_code == 404:
        print("error: /v2/ not found")
        return
    else:
        print(f"error: unknown http_code {http_code}")
        return

    print("|> get manifest.list(index.json) digest")
    resp = session.head(
        f"https://{host}/v2/{img}/manifests/{tag}",
        headers={"Accept": "application/vnd.docker.distribution.manifest.list.v2+json"},
    )
    if resp.status_code != 200:
        print(f"error: failed to get manifest.list digest {resp.status_code}")
        return
    digest = resp.headers.get("Docker-Content-Digest")
    if not digest:
        print("error: Docker-Content-Digest header not found")
        return
    print(f"Manifest Digest: {digest}")

    print("|> get manifest.list(index.json)")
    resp = session.get(
        f"https://{host}/v2/{img}/manifests/{digest}",
        headers={"Accept": "application/vnd.docker.distribution.manifest.list.v2+json"},
    )
    resp.raise_for_status()
    content_type = resp.headers['Content-Type']
    manifest_text = resp.text
    manifest = json.loads(manifest_text)
    if content_type == 'application/vnd.docker.distribution.manifest.v2+json':
        print("got manifest, skip manifest.list")
    elif content_type == "application/vnd.docker.distribution.manifest.list.v2+json":
        if "manifests" not in manifest:
            print("error: manifests not found in manifest.list")
            print(manifest_text)
            return
        with open(f"{save_path}/index.json", "w") as f:
            f.write(manifest_text)

        # select amd64_manifest
        amd64_manifest = None
        for m in manifest["manifests"]:
            if m.get("platform", {}).get("architecture") == "amd64":
                amd64_manifest = m
                break
        if not amd64_manifest:
            print("error: amd64_manifest not fould")
            print(manifest_text)
            return
        amd64_digest = amd64_manifest["digest"]
        print(f"AMD64 Manifest Digest: {amd64_digest}")

        print("|> get manifest blobs")
        resp = session.get(
            f"https://{host}/v2/{img}/manifests/{amd64_digest}",
            headers={"Accept": "application/vnd.docker.distribution.manifest.v2+json"},
        )
        resp.raise_for_status()
        manifest_text = resp.text
        manifest = json.loads(manifest_text)
        digest = amd64_digest
    else:
        print(f"error: unknown Content-Type {content_type}")
        return

    with open(f"{blobs_path}/{digest.replace('sha256:', '')}", "w") as f:
        f.write(manifest_text)

    # https://github.com/openshift/docker-distribution/blob/main/docs/spec/manifest-v2-2.md
    if manifest["schemaVersion"] != 2:
        print("required application/vnd.docker.distribution.manifest.v2+json")
        return

    print("|> get config blobs")
    config_digest = manifest["config"]["digest"]
    config_path = f"{blobs_path}/{config_digest.replace('sha256:', '')}"
    resp = session.get(f"https://{host}/v2/{img}/blobs/{config_digest}")
    resp.raise_for_status()
    config_text = resp.text
    with open(config_path, "w") as f:
        f.write(config_text)

    print("|> parse layer digests")
    layer_sources = {layer["digest"]: layer for layer in manifest["layers"]}
    layer_digests = [layer["digest"] for layer in manifest["layers"]]

    print("all layer digests:")
    for layer in layer_digests:
        print(layer)

    blobs_url_prefix = f"https://{host}/v2/{img}/blobs/"
    download_list = ""
    for layer in layer_digests:
        download_url = blobs_url_prefix + layer
        resp = session.head(download_url)
        if resp.status_code != 307:
            print("HEAD download_url status_code is not 307 redirect")
            return
        download_url = resp.headers.get("Location", "")
        download_list += download_url + "\n"
        download_list += "  out=" + layer.replace("sha256:", "") + "\n"

    manifest_json = [
        {
            "Config": config_path.replace(f"{save_path}/", ""),
            "RepoTags": [repo_tag],
            "Layers": [f"blobs/{x.replace(':', '/')}" for x in layer_digests],
            "LayerSources": layer_sources,
        }
    ]
    repo = {repo_tag.split(":")[0]: {tag: layer_digests[-1][len("sha256:") :]}}

    with open(f"{save_path}.txt", "w") as f:
        f.write(download_list)
    with open(f"{save_path}/oci-layout", "w") as f:
        f.write('{"imageLayoutVersion": "1.0.0"}')
    with open(f"{save_path}/repositories", "w") as f:
        f.write(json.dumps(repo))
    with open(f"{save_path}/manifest.json", "w") as f:
        f.write(json.dumps(manifest_json))

    print("<| Download Image |>")
    print(
        f'aria2c -x16 -s16 -j16 -k1M --auto-file-renaming=false --file-allocation=falloc --dir={blobs_path} --header="{DOCKER_USER_AGENT}" -i {save_path}.txt'
    )

    print("<| Pack Image |>")
    print(f"tar -cvf {save_path}.tar -C {save_path} .")

    print("<| Load Image |>")
    print(f"docker load -i {save_path}.tar")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "repo",
        help="docker repo like debian, debian:bookworm, nvcr.io/nvidia/cuda:latest",
    )
    args = parser.parse_args()
    docker_pull(args.repo)
